import sys
import argparse
import logging
import faulthandler
from pathlib import Path
from multiprocessing import Array

from nodes.Node import Node
from nodes.SourceNode import SourceNode, VideoSourceNode, RtspSourceNode
from nodes.SinkNode import VideoSinkNode, TestSinkNode, PlayerSinkNode, RtmpSinkNode
from nodes.FFVideoWriterNode import FFVideoWriterNode
from nodes.YOLONode import YOLONode
from nodes.SORTTrackNode import SORTTrackNode
from nodes.FlaskSinkNode import FlaskSinkNode
from nodes.VideoListSourceNode import VideoListSourceNode
from nodes.utils.shm import *
from FishVideoListSourceNode import FishVideoListSourceNode
from FishSORTTrackNode import FishSORTTrackNode
from FishFFVideoWriterNode import FishFFVideoWriterNode
from FishYOLONode import FishYOLONode
from FishBatchYOLONode import FishBatchYOLONode
from FishFilterNode import FishFilterNode
from utils import *
#from FlaskFetcher import fishFetcher

level = logging.ERROR
logging.basicConfig(level=level)

faulthandler.enable()

#def createAppDict(key,nc):
#    dict ={}
#    for i in range(nc):
#        dict[i]=0
#
#    CreateAppSharedMemory(key)
#
#    StoreAppDict(key, dict)

def getParams():
    # Create an ArgumentParser object
    parser = argparse.ArgumentParser(description="A simple argument parser")

    # Define the arguments you expect
    parser.add_argument('--input', type=str, help='A video folder or a video list file.', default="../input")
    parser.add_argument('--output', type=str, help='A output folder.', default="../output")
#    parser.add_argument('--appkey', type=str, help='A unique string for communication', default="/FisherKey")
    parser.add_argument('--port', type=int, help="Video port.", default=8042)

    # Parse the arguments
    args = parser.parse_args()

    return args

def main():
    #CreateSharedMemory(SHM_VIDEO_NAME);
    args = getParams()

    strinput = args.input
    stroutput = args.output
    fport = args.port
#    appkey = args.appkey

#    createAppDict(appkey, 5)
#    dict = LoadAppDict(appkey)
#    print('AppDict:', dict)

    globalStatus = Array('i', [0, 0, 0, 0, 0])

    processes = []

    source = FishVideoListSourceNode()
    source.listVideos(strinput)
    source.setId(0)
    #source.setAppKey(appkey)
    source.setStatusArray(globalStatus)
    source.setStatus(1)
    source.saveVideoList(stroutput + "/videolist.txt")
    source.setSourceCloseAction(source.closeAction)

    yolo = FishBatchYOLONode()
    yolo.setId(1)
    #yolo.setAppKey(appkey)
    yolo.setStatusArray(globalStatus)
    yolo.setBatchSize(50)
    yolo.updatePoolSize(8000)
    yolo.setModelFile("../FishZHX20241204A.pt")

    sort = FishSORTTrackNode()
    sort.setId(2)
    #sort.setAppKey(appkey)
    sort.setStatusArray(globalStatus)
    sort.updatePoolSize(2000)

    filter = FishFilterNode()
    filter.setId(3)
    filter.setStatusArray(globalStatus)
    #filter.setAppKey(appkey)
    filter.updatePoolSize(2000)

    #fsink = FlaskSinkNode()
    #fsink.updatePoolSize(500)
    #fsink.setPort(fport)

    vsink = FishFFVideoWriterNode()
    vsink.setId(4)
    #vsink.setAppKey(appkey)
    vsink.setStatusArray(globalStatus)
    vsink.updatePoolSize(7000)
    vsink.setFrameSize((960,540))
    vsink.setVideoPath(stroutput)
    vsink.setLogFileName(stroutput + "/fish.txt")
    vsink.openLog()


    #fsink.setParent(filter)
    vsink.setParent(filter)
    filter.setParent(sort)
    sort.setParent(yolo)
    yolo.setParent(source)

    #fsink.start()
    vsink.start()
    processes.append(vsink)
    filter.start()
    processes.append(filter)
    sort.start()
    processes.append(sort)
    yolo.start()
    processes.append(yolo)
    source.start()
    processes.append(source)

    while(True):
        allexited = True
        for process in processes:
            process.join()
            if process.is_alive():
                process.terminate()
                print(f"XXXXXXXXXXXXXXXXXXX+++++++++ Terminated process {process.name} +++++++XXXXXXXXXXXxXX")
            else:
                print(f"EEEEEEEEEEEEEEEEEEE+++++++++ Process {process.name} +++++++EEEEEEEEEEEE")

if __name__ == "__main__":
    main()
